package encrypt

import (
	"crypto"
	"crypto/md5"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"crypto/x509"
	"encoding/base64"
	"fmt"
	"hash"
	"strings"
)


// 加密数据
func RSAEncrypt(publicKeyStr string, plaintext string) (string, error) {
	publicKey, err := parseRSAPublicKey(publicKeyStr)
	if err != nil {
		return "", err
	}

	ciphertext, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, []byte(plaintext))
	if err != nil {
		return "", err
	}

	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// 解密数据
func RSADecrypt(privateKeyStr string, ciphertext string) (string, error) {
	privateKey, err := parseRSAPrivateKey(privateKeyStr)
	if err != nil {
		return "", err
	}

	decodedCiphertext, err := base64.StdEncoding.DecodeString(ciphertext)
	if err != nil {
		return "", err
	}

	plaintext, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, decodedCiphertext)
	if err != nil {
		return "", err
	}

	return string(plaintext), nil
}

// 对数据进行签名
func RSASign(privateKeyStr string, data string, hashAlgorithm crypto.Hash) (string, error) {
	privateKey, err := parseRSAPrivateKey(privateKeyStr)
	if err != nil {
		return "", err
	}

	hashed := hashData(data, hashAlgorithm)

	signature, err := rsa.SignPKCS1v15(rand.Reader, privateKey, hashAlgorithm, hashed)
	if err != nil {
		return "", err
	}

	return base64.StdEncoding.EncodeToString(signature), nil
}

// 验证RSA签名
func RSASignVerify(publicKeyStr string, data string, signature string, hashAlgorithm crypto.Hash) bool {
	publicKey, err := parseRSAPublicKey(publicKeyStr)
	if err != nil {
		return false
	}

	hashed := hashData(data, hashAlgorithm)
	sign, err := base64.StdEncoding.DecodeString(signature)
	if err != nil {
		return false
	}
	err = rsa.VerifyPKCS1v15(publicKey, hashAlgorithm, hashed, sign)
	if err != nil {
		return false
	}

	return true
}

// 解析RSA公钥
func parseRSAPublicKey(publicKeyStr string) (*rsa.PublicKey, error) {
	// 移除字符串中的换行符和空格
	publicKeyStr = strings.ReplaceAll(publicKeyStr, "-----BEGIN PUBLIC KEY-----", "")
	publicKeyStr = strings.ReplaceAll(publicKeyStr, "-----END PUBLIC KEY-----", "")
	publicKeyStr = strings.ReplaceAll(publicKeyStr, "\n", "")
	publicKeyStr = strings.ReplaceAll(publicKeyStr, "\t", "")
	publicKeyStr = strings.ReplaceAll(publicKeyStr, " ", "")

	// 解码Base64
	decodedKey, err := base64.StdEncoding.DecodeString(publicKeyStr)
	if err != nil {
		return nil, err
	}

	// 解析公钥
	publicKeyInterface, err := x509.ParsePKIXPublicKey(decodedKey)
	if err != nil {
		return nil, err
	}

	publicKey, ok := publicKeyInterface.(*rsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("无法解析RSA公钥")
	}

	return publicKey, nil
}

// 解析RSA私钥
func parseRSAPrivateKey(privateKeyStr string) (*rsa.PrivateKey, error) {
	// 移除字符串中的换行符和空格
	privateKeyStr = strings.ReplaceAll(privateKeyStr, "-----BEGIN RSA PRIVATE KEY-----", "")
	privateKeyStr = strings.ReplaceAll(privateKeyStr, "-----END RSA PRIVATE KEY-----", "")
	privateKeyStr = strings.ReplaceAll(privateKeyStr, "\n", "")
	privateKeyStr = strings.ReplaceAll(privateKeyStr, " ", "")
	// 解码Base64
	decodedKey, err := base64.StdEncoding.DecodeString(privateKeyStr)
	if err != nil {
		return nil, err
	}
	/* 	privateKeyBlock, _ := pem.Decode([]byte(privateKeyStr))
	   	if privateKeyBlock == nil {
	   		return nil, fmt.Errorf("无法解析私钥")
	   	} */

	privateKey, err := x509.ParsePKCS1PrivateKey(decodedKey)
	if err != nil {
		return nil, err
	}

	return privateKey, nil
}

// 计算数据的哈希值
func hashData(data string, hashAlgorithm crypto.Hash) []byte {
	var h hash.Hash

	switch hashAlgorithm {
	case crypto.MD5:
		h = md5.New()
	case crypto.SHA1:
		h = sha1.New()
	case crypto.SHA224:
		h = sha256.New224()
	case crypto.SHA256:
		h = sha256.New()
	case crypto.SHA384:
		h = sha512.New384()
	case crypto.SHA512:
		h = sha512.New()
	default:
		h = sha256.New()
	}

	h.Write([]byte(data))
	return h.Sum(nil)
}
